package hw2.memdir;

import hw2.FileIO;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.ObjectInputStream;
import java.security.InvalidKeyException;
import java.util.concurrent.ConcurrentHashMap;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

/**
 * 
 * The interface to query the encrypted data
 * 
 */
public class ListQuery
{

    // maximum bad password attempts
    private static final int MAX_ATTEMPTS = 3;

    // lastName ---> Email
    static ConcurrentHashMap<String, String> memInfo = new ConcurrentHashMap<String, String>();

    static Cipher aesEnc;

    // static Cipher aesDec;
    static byte[] iv;

    // the key to be generated by existing IV, salt and provided password
    static SecretKey aesKey = null;

    static Directory dir;

    /**
     * @param args
     * @throws IOException
     * @throws InvalidKeyException
     */
    public static void main( String[] args ) throws Exception {
        BufferedReader in = new BufferedReader( new InputStreamReader(
                System.in ) );

        String s = "";

        // continue trying until correct password is entered
        Boolean getPassword = true;

        // load the encrypted directory
        FileInputStream fis = new FileInputStream( new File( "data/dir.enc" ) );
        ObjectInputStream ois = new ObjectInputStream( fis );
        dir = (Directory) ois.readObject();

        // load the IV used for last names encryption
        iv = FileIO.getBytesFromFile( "data/iv.bin" );

        // track password attempts
        int count = 0;

        // prompt for passwords
        while( getPassword && count < MAX_ATTEMPTS ) {
            s = "";
            System.out.print( "password >> " );
            s = in.readLine();
            count++;

            // generate an AES key with the provided password
            aesKey = KeyGen.getAESkey( s );

            // create the encryption cipher for encrypting queries
            aesEnc = Cipher.getInstance( "AES/CBC/PKCS5Padding" );
            aesEnc.init( Cipher.ENCRYPT_MODE, aesKey,
                    new IvParameterSpec( dir.getNameIV() ) );

            // check if the key works (i.e. is the password correct
            if( isKeyValid( aesKey ) ) {
                getPassword = false;
            } else {
                if( count < MAX_ATTEMPTS ) {
                    System.out.println( "Incorrect password. Try again." );
                } else {
                    System.out.println( "Too many bad attempts. Goodbye." );
                    System.exit( 1 );
                }
            }

        }

        // correct password has been entered and queries can proceed
        runQueries();

    }

    /**
     * 
     * @param proposedAESkey
     *            a key to check for correctness
     * @return
     */
    private static boolean isKeyValid( SecretKey proposedAESkey ) {
        boolean correctKey = false;

        // an entry "1235" with the value of "correct" is in the member list
        // if the key is correct, the lookup should be successful
        try {
            if( query( "12345" ).equalsIgnoreCase( "correct" ) ) {
                correctKey = true;
            }
        } catch( Exception e ) {
            e.printStackTrace();
        }
        return correctKey;
    }

    /**
     * 
     * @param str
     *            a query string
     * @return
     * @throws Exception
     */
    private static String query( String str ) throws Exception {
        // encrypt the query
        byte[] c = aesEnc.doFinal( str.getBytes() );

        // query with encrypted name
        byte[] respiv = dir.getIV( c );
        if( respiv != null ) {
            byte[] respEmail = dir.getEncryptedEmail( c );
            Cipher aesDec = Cipher.getInstance( "AES/CBC/PKCS5Padding" );
            aesDec.init( Cipher.DECRYPT_MODE, aesKey, new IvParameterSpec(
                    respiv ) );
            String email = new String( aesDec.doFinal( respEmail ) );
            return email;
        } else {
            return "Not Found.";
        }
    }

    /**
     * Prompts for input from user, checks for quit and says goodbye
     * 
     * @throws Exception
     */
    private static void runQueries() throws Exception {
        BufferedReader in = new BufferedReader( new InputStreamReader(
                System.in ) );
        String s = "";
        System.out.println( "\nWelcome.  Enter \"quit\" to exit." );
        while( true ) {
            System.out.print( "Last Name >> " );
            s = in.readLine().toLowerCase();
            if( "quit".equalsIgnoreCase( s ) )
                break; // end program
            // Enc s using shared key
            // Send to query program
            System.out.println( query( s ) );
        }
        System.out.println( "Goodbye." );
    }

}
